package com.snail.handler;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class OnlineUserWebSocketHandler extends TextWebSocketHandler {

    private final RedisTemplate redisTemplate;

    // 存储所有WebSocket会话
    private final ConcurrentHashMap<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    private final ObjectMapper objectMapper = new ObjectMapper();

    public OnlineUserWebSocketHandler(RedisTemplate redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    /**
     * 连接建立后
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.put(session.getId(), session);
        log.info("WebSocket连接建立，会话ID: {}", session.getId());
        log.info("获取当前登录用户信息");

        // 广播更新在线用户数
        broadcastOnlineStats();

        log.info("WebSocket连接建立完成，会话ID: {}", session.getId());
    }

    /**
     * 处理收到的消息
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        log.debug("收到WebSocket消息: {}", payload);

        // 处理心跳包
        if ("ping".equals(payload)) {
            session.sendMessage(new TextMessage("pong"));
            // 在收到ping时也发送当前在线统计信息
            sendCurrentStatsToSession(session);
        }
    }

    /**
     * 连接关闭后
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        sessions.remove(session.getId());

        // 广播更新在线用户数
        broadcastOnlineStats();

        log.info("WebSocket连接关闭，会话ID: {}，状态: {}", session.getId(), status);
    }

    /**
     * 传输错误处理
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输错误，会话ID: {}", session.getId(), exception);
        sessions.remove(session.getId());
        broadcastOnlineStats();
    }

    /**
     * 广播在线统计信息
     */
    private void broadcastOnlineStats() throws IOException {
        OnlineStats stats = new OnlineStats();
        stats.setOnlineCount(getOnlineUserCountFromRedis());
        stats.setTimestamp(System.currentTimeMillis());

        log.debug("广播在线统计信息，当前在线用户数: {}", getOnlineUserCountFromRedis());

        String message = objectMapper.writeValueAsString(stats);

        // 向所有连接的客户端广播消息
        for (WebSocketSession session : sessions.values()) {
            if (session.isOpen()) {
                try {
                    session.sendMessage(new TextMessage(message));
                    log.debug("向会话 {} 发送统计信息: {}", session.getId(), message);
                } catch (IOException e) {
                    log.error("发送WebSocket消息失败", e);
                }
            }
        }
    }

    /**
     * 向单个会话发送当前统计信息
     */
    private void sendCurrentStatsToSession(WebSocketSession session) {
        try {
            OnlineStats stats = new OnlineStats();
            stats.setOnlineCount(getOnlineUserCountFromRedis());
            stats.setTimestamp(System.currentTimeMillis());

            String message = objectMapper.writeValueAsString(stats);
            if (session.isOpen()) {
                session.sendMessage(new TextMessage(message));
                log.debug("向会话 {} 发送统计信息: {}", session.getId(), message);
            }
        } catch (IOException e) {
            log.error("向会话 {} 发送统计信息失败", session.getId(), e);
        }
    }

    /**
     * 从Redis获取当前在线用户数
     */
    private int getOnlineUserCountFromRedis() {
        try {
            Set<String> keys = redisTemplate.keys("sa-token:login:sess*");
            return keys != null ? keys.size() : 0;
        } catch (Exception e) {
            log.error("从Redis获取在线用户数失败", e);
            return 0;
        }
    }

    /**
     * 获取当前在线用户数
     */
    public int getOnlineUserCount() {
        log.debug("获取当前在线用户数");
        return getOnlineUserCountFromRedis();
    }
    
    /**
     * 获取当前WebSocket连接数
     */
    public ConcurrentHashMap<String, WebSocketSession> getSessions() {
        return sessions;
    }

    /**
     * 统计数据类
     */
    public static class OnlineStats {
        private int onlineCount;
        private long timestamp;

        // Getters and Setters
        public int getOnlineCount() {
            return onlineCount;
        }

        public void setOnlineCount(int onlineCount) {
            this.onlineCount = onlineCount;
        }

        public long getTimestamp() {
            return timestamp;
        }

        public void setTimestamp(long timestamp) {
            this.timestamp = timestamp;
        }
    }

    /**
     * 在线用户信息类
     */

}